import os, sys, time, argparse
import cv2
from paddleocr import PaddleOCR

from loguru import logger

HOME = os.path.expanduser('~')
MODEL_DIR = os.path.join(HOME, ".paddlex\\official_models")

def ocr_init():
    # ocr = PaddleOCR(
    #     use_doc_orientation_classify=False,
    #     use_doc_unwarping=False,
    #     use_textline_orientation=False) # 文本检测+文本识别
    # ocr = PaddleOCR(use_doc_orientation_classify=True, use_doc_unwarping=True) # 文本图像预处理+文本检测+方向分类+文本识别
    # ocr = PaddleOCR(use_doc_orientation_classify=False, use_doc_unwarping=False) # 文本检测+文本行方向分类+文本识别
    # bgn = time.time()

    logger.info("bgn")
    ocr = PaddleOCR(lang="ch", # device="gpu:0",
        text_detection_model_name="PP-OCRv5_mobile_det",
        text_detection_model_dir=os.path.join(MODEL_DIR, "PP-OCRv5_mobile_det"),
        text_recognition_model_name="PP-OCRv5_mobile_rec",
        text_recognition_model_dir=os.path.join(MODEL_DIR, "PP-OCRv5_mobile_rec"),
        use_doc_orientation_classify=False,
        use_doc_unwarping=False,
        use_textline_orientation=False) # 更换 PP-OCRv5_mobile 模型

    logger.info("loaded model")
    return ocr

def ocr_image(ocr, img):
    logger.info("predict bgn")
    result = ocr.predict(img)
    logger.info("predict end")
    return result

def threshold_image(pic):
    logger.info("threshold {}".format(pic))
    img = cv2.imread(pic)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY);
    # 二值化
    # ret, th1 = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY);
    ret, th2 = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
    th_img = th2
    if False:
        # 显示图像
        # cv2.imshow("th1", th1);
        cv2.imshow("th2", th2);
        # 进程不结束，一直保持显示状态
        cv2.waitKey(0);
        #销毁所有窗口
        cv2.destroyAllWindows();
    return th_img

def test():
    ocr = ocr_init()
    for pic in ("test/test.jpg", "test/test1.png"):
        if False:
            th_img = threshold_image(pic)
            img = cv2.cvtColor(th_img, cv2.COLOR_GRAY2BGR)
        else: img = cv2.imread(pic)
        if True:
            result = ocr_image(ocr, img)
            for res in result:
                res.print()
                res.save_to_img("output")
                res.save_to_json("output")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # parser.add_argument('action', choices=['cli', 'gui'])
    parser.add_argument('-c', '--config', type=str, help='config file')
    parser.add_argument('-v', '--version', type=str, help='version')
    args = parser.parse_args()

    test()
